import cv2
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks


class ImageMagic(object):
    def __init__(self):
        self.model_path = {
            "style_cartoon": 'damo/cv_unet_person-image-cartoon_compound-models',
            "style_3d": 'damo/cv_unet_person-image-cartoon-3d_compound-models',
            "style_hand_drawn": 'damo/cv_unet_person-image-cartoon-handdrawn_compound-models',
            "style_sketch": 'damo/cv_unet_person-image-cartoon-sketch_compound-models',
            "style_art": 'damo/cv_unet_person-image-cartoon-artstyle_compound-models',
            "style_transfer": 'damo/cv_aams_style-transfer_damo'
        }

    @staticmethod
    def save(pipeline_obj, img_des):
        cv2.imwrite(img_des, pipeline_obj[OutputKeys.OUTPUT_IMG])
        print('finished!')

    def stylization(self, img_src, img_des, style="cartoon"):
        img_cartoon = pipeline(Tasks.image_portrait_stylization, model=self.model_path[f"style_{style}"])
        result = img_cartoon(img_src)
        self.save(result, img_des)

    def style_transfer(self, img_src, img_style, img_des):
        style_transfer = pipeline(Tasks.image_style_transfer, model_id=self.model_path["style_transfer"])
        result = style_transfer(dict(content=img_src, style=img_style))
        self.save(result, img_des)


if __name__ == '__main__':
    im = ImageMagic()

    image_1 = "share.jpg"
    style_img = 'https://modelscope.oss-cn-beijing.aliyuncs.com/demo/image-style-transfer/style_transfer_style.jpg'

    # im.style_transfer(image_1, style_img, "result2.png")

    im.stylization(image_1, "result3.png")
